﻿using System.Collections.Generic;
using System.Runtime.InteropServices;
using Unity.Burst;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Jobs;
using Unity.Mathematics;
using UnityEngine;
using UnityEngine.Rendering;
using Random = Unity.Mathematics.Random;

namespace Learning_GPUInstancing_6_Simple4_2
{
    /// <summary>
    /// 用通用的 NativeArray<byte> 来承载所有类型的数据
    /// </summary>
    public unsafe class Simple4GPUInstancing_2 : MonoBehaviour
    {
        /// <summary>
        /// GPU Instancing 一个批次可提交的数量
        /// 注意：这里设置得越大，DC会越少，但是特殊情况下浪费的带宽会越多（比如只提交10个时，一个批次还是会提交512个数据给GPU，但是有效数据只有10个）
        /// </summary>
        private const int BATCH = 256;
        
        private const float Range = 100000f;
        private static readonly Bounds BOUNDS = new Bounds(Vector3.zero, new Vector3(Range, Range, Range));

        private static readonly int Float4x4_Size = Marshal.SizeOf<float4x4>();

        private static readonly int TransID = Shader.PropertyToID("transformBuf");
        
        public int count = 100000;
        public Mesh mesh;
        public Material material;

        private NativeArray<SampleData> _dataArr;
        private List<BatchInfo> _batchInfoList = new List<BatchInfo>();
        
        void Start()
        {
            _dataArr = new NativeArray<SampleData>(count, Allocator.Persistent);
            
            var job = new InitJob();
            job.array = _dataArr;
            var jobHandle = job.Schedule(count, 64);
            jobHandle.Complete();
        }
        
        void Update()
        {
            var transArr = new NativeArray<float4x4>(count, Allocator.TempJob);
            
            var job = new UpdateJob();
            job.dataArr = _dataArr;
            job.transArr = transArr;
            job.deltaTime = Time.deltaTime;
            var jobHandle = job.Schedule(count, 64);
            jobHandle.Complete();
            
            var batchCount = (int) math.ceil((float) count / (float) BATCH);
            for (int i = 0; i < batchCount; i++)
            {
                var startIndex = i * BATCH;
                var len = BATCH;
                if (i == batchCount - 1)
                {
                    len = count - (BATCH * i);
                }

                if (_batchInfoList.Count <= i)
                {
                    var info = new BatchInfo();
                    
                    info.block = new MaterialPropertyBlock();
                    
                    info.args = new uint[5] {0, 0, 0, 0, 0};
                    info.argsBuf = new GraphicsBuffer(GraphicsBuffer.Target.IndirectArguments, 1, 5 * sizeof(uint));
                    
                    info.transData = new NativeArray<byte>(BATCH * Float4x4_Size, Allocator.Persistent);
                    info.transBuf = new GraphicsBuffer(GraphicsBuffer.Target.Structured, BATCH, Marshal.SizeOf(typeof(float4x4)));
                    
                    _batchInfoList.Add(info);
                }
                var batchInfo = _batchInfoList[i];
                batchInfo.block.Clear();

                var destPtr = (float4x4*) batchInfo.transData.GetUnsafePtr();
                var sourcePtr = (float4x4*) transArr.GetUnsafePtr();
                
                UnsafeUtility.MemCpy(destPtr, sourcePtr + startIndex, len * Float4x4_Size);
                batchInfo.transBuf.SetData(batchInfo.transData);
                batchInfo.block.SetBuffer(TransID, batchInfo.transBuf);
                
                batchInfo.args[0] = mesh.GetIndexCount(0); // 每个实例的索引数
                batchInfo.args[1] = (uint) len; // 实例数
                batchInfo.args[2] = mesh.GetIndexStart(0); // 起始索引位置
                batchInfo.args[3] = mesh.GetBaseVertex(0); // 基顶点位置
                batchInfo.args[4] = 0; // 起始实例位置
                batchInfo.argsBuf.SetData(batchInfo.args);

                // 绘制
                Graphics.DrawMeshInstancedIndirect(mesh, 0, material, BOUNDS, batchInfo.argsBuf,
                    0, batchInfo.block, ShadowCastingMode.On, true, LayerMask.NameToLayer("Character"));
            }
            
            transArr.Dispose();
        }

        private void OnDestroy()
        {
            _dataArr.Dispose();
        }
    }
    
    [BurstCompile]
    public struct InitJob : IJobParallelFor
    {
        public NativeArray<SampleData> array;
        
        public unsafe void Execute(int index)
        {
            var random = Random.CreateFromIndex((uint) index);

            var ptr = (SampleData*) array.GetUnsafePtr();
            
            var data = ptr + index;
            
            data->pos = new float3(random.NextFloat(-50f, 50f), random.NextFloat(-50f, 50f), random.NextFloat(-50f, 50f));
            data->rotation = quaternion.Euler(random.NextFloat(360f), random.NextFloat(360f), random.NextFloat(360f));
            var scale = random.NextFloat(0.5f, 1.5f);
            data->scale = new float3(scale, scale, scale);

            data->forward = math.normalize(math.forward(data->rotation));
            data->rotateAdd = new float3(random.NextFloat(10f), random.NextFloat(10f), random.NextFloat(10f));
        }
    }
    
    [BurstCompile]
    public struct UpdateJob : IJobParallelFor
    {
        public NativeArray<SampleData> dataArr;
        public NativeArray<float4x4> transArr;
        public float deltaTime;
        
        public unsafe void Execute(int index)
        {
            var dataPtr = (SampleData*) dataArr.GetUnsafePtr();
            var transPtr = (float4x4*) transArr.GetUnsafePtr();
            
            var data = dataPtr + index;

            data->pos += data->forward * deltaTime;
            data->rotation = math.mul(data->rotation, quaternion.Euler(data->rotateAdd * deltaTime));
            
            var trans = float4x4.TRS(data->pos, data->rotation, data->scale);
            *(transPtr + index) = trans;
        }
    }
    
    public struct SampleData
    {
        public float3 pos;
        public quaternion rotation;
        public float3 scale;

        public float3 forward; // 1秒向前移动的位移
        public float3 rotateAdd; // 1秒旋转的增量
    }

    public class BatchInfo
    {
        public MaterialPropertyBlock block;
        
        public uint[] args;
        public GraphicsBuffer argsBuf;
        
        public NativeArray<byte> transData;
        public GraphicsBuffer transBuf;
    }
}
